package ccnl.demo.algo;

import ccnl.demo.JrddTools;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.ml.regression.RandomForestRegressor;
import org.apache.spark.mllib.classification.LogisticRegressionModel;
import org.apache.spark.mllib.classification.LogisticRegressionWithLBFGS;
import org.apache.spark.mllib.regression.LabeledPoint;
import org.apache.spark.mllib.tree.RandomForest;
import org.apache.spark.sql.DataFrame;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SQLContext;
import scala.Tuple2;

/**
 * Created by wong on 16/5/11.
 */
public class RF {
    JavaSparkContext javaSparkContext;
    SQLContext sqlContext;
    public RF(JavaSparkContext jsc, SQLContext sqlc){
        this.javaSparkContext = jsc;
        this.sqlContext = sqlc;
    }


    public void runJob(String datasetPath) {
        DataFrame df = sqlContext.read().parquet(datasetPath);
        JavaRDD<Row> jrdd = df.toJavaRDD();
        JavaRDD<Row> jrdd_train = JrddTools.getTrain(jrdd);
        JavaRDD<Row> jrdd_test = JrddTools.getTest(jrdd);

        JavaRDD<LabeledPoint> jrdd_lp_train = JrddTools.convertRow2Lp(jrdd_train);
        JavaRDD<LabeledPoint> jrdd_lp_test = JrddTools.convertRow2Lp(jrdd_test);


        JavaRDD<Tuple2<Object, Object>> predictionAndLabels = runGiven(jrdd_lp_train, jrdd_lp_test);
//        printMetricsLog("no distinct trainSet / 1:160.5", predictionAndLabels, jrdd_lp_train, jrdd_lp_test);

    }

    public JavaRDD<Tuple2<Object, Object>> runGiven(JavaRDD<LabeledPoint> train, JavaRDD<LabeledPoint> test) {
        RandomForestRegressor rf = new RandomForestRegressor().setLabelCol("C1").setFeaturesCol("features");


        train.cache();
        final LogisticRegressionModel model = new LogisticRegressionWithLBFGS()
                .run(train.rdd());
        model.clearThreshold();
        JavaRDD<Tuple2<Object, Object>> predictionAndLabels = test.map(v1 -> {
            Double prediction = model.predict(v1.features());
            return new Tuple2<>(prediction, v1.label());
        });
        return predictionAndLabels;
    }
}
